package com.kaigle.ai.controller;

import com.kaigle.ai.tools.EmailTool;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.memory.InMemoryChatMemoryRepository;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.deepseek.DeepSeekChatModel;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.HashMap;

@RestController
@RequestMapping("/ai")
public class AiChatController {
    @Resource
    private DeepSeekChatModel deepSeekChatModel;
    @Resource
    private EmailTool emailTool;


    @GetMapping("/deepSeek/call")
    public String deepSeek(@RequestParam("msg") String msg) {
        Prompt prompt = new Prompt(msg);
        ChatResponse chatResponse = deepSeekChatModel.call(prompt);
        return chatResponse.getResult().getOutput().getText();
    }

    @GetMapping("/deepSeek/stream")
    public String deepSeekStream(@RequestParam("msg") String msg) {
        Prompt prompt = new Prompt(msg);
        Flux<ChatResponse> stream = deepSeekChatModel.stream(prompt);
        // 阻塞输出
        stream.toIterable().forEach(System.out::print);
        return "ok";
    }

    @GetMapping("/chatClient/call")
    public String chatClient(@RequestParam("msg") String msg) {
        ChatClient chatClient = ChatClient.builder(deepSeekChatModel).build();
        Prompt prompt = new Prompt(msg);
        ChatResponse chatResponse = chatClient.prompt(prompt)
                .call().chatResponse();
        return chatResponse.getResult().getOutput().getText();
    }

    @Resource
    private ChatMemoryRepository chatMemoryRepository;

    @GetMapping(value = "/chatClient/stream", produces = "text/event-stream")
    public Flux<String> chatClientStream(@RequestParam("userName") String userName,@RequestParam("msg") String msg) {
        ChatClient chatClient = ChatClient.builder(deepSeekChatModel).build();
        Prompt prompt = new Prompt(msg);
        ChatClient.StreamResponseSpec stream = chatClient.prompt(prompt)
                .toolContext(new HashMap<>())
                .advisors(
                        // 对话记忆
                        MessageChatMemoryAdvisor.builder(
                                // 基于内存 InMemoryChatMemoryRepository 的对话记忆
                                MessageWindowChatMemory.builder().chatMemoryRepository(chatMemoryRepository).build()).build(),
                        // 日志
                        new SimpleLoggerAdvisor()
                     )
                // 多用户对话记忆隔离
                .advisors(advisorSpec -> advisorSpec.param(ChatMemory.CONVERSATION_ID,userName))
                // 可调用工具
                .tools(emailTool)
                .stream();
        return stream.content();
    }

}
